# Apriori算法代码
# 参考 https://github.com/asaini/Apriori

def load_dataset():
    # 加载示例数据集
    dataset = [['Milk', 'Onion', 'Nutmeg', 'Kidney Beans', 'Eggs', 'Yogurt'],
               ['Dill', 'Onion', 'Nutmeg', 'Kidney Beans', 'Eggs', 'Yogurt'],
               ['Milk', 'Apple', 'Kidney Beans', 'Eggs'],
               ['Milk', 'Unicorn', 'Corn', 'Kidney Beans', 'Yogurt'],
               ['Corn', 'Onion', 'Onion', 'Kidney Beans', 'Ice cream', 'Eggs']]

    return dataset

def create_candidate_set(dataset):
    # 生成候选项集
    candidate_set = []
    for transaction in dataset:
        for item in transaction:
            if not [item] in candidate_set:
                candidate_set.append([item])

    candidate_set.sort()
    return list(map(frozenset, candidate_set))

def scan_dataset(dataset, candidate_set, min_support):
    # 计算候选项集的支持度
    item_count = {}
    for transaction in dataset:
        for candidate in candidate_set:
            if candidate.issubset(transaction):
                item_count[candidate] = item_count.get(candidate, 0) + 1

    num_items = float(len(dataset))
    freq_items = []
    support_data = {}
    for key in item_count:
        support = item_count[key] / num_items
        if support >= min_support:
            freq_items.insert(0, key)
        support_data[key] = support

    return freq_items, support_data

def apriori_gen(freq_items, k):
    # 生成频繁项集的候选项集
    next_candidates = []
    num_freq_items = len(freq_items)
    for i in range(num_freq_items):
        for j in range(i + 1, num_freq_items):
            L1 = list(freq_items[i])[:k - 2]
            L2 = list(freq_items[j])[:k - 2]
            L1.sort()
            L2.sort()
            if L1 == L2:
                next_candidates.append(freq_items[i] | freq_items[j])

    return next_candidates

def apriori(dataset, min_support=0.5):
    dataset = list(map(set, dataset))
    candidate_set = create_candidate_set(dataset)
    freq_items1, support_data = scan_dataset(dataset, candidate_set, min_support)
    freq_items = [freq_items1]
    k = 2

    while len(freq_items[k - 2]) > 0:
        next_candidates = apriori_gen(freq_items[k - 2], k)
        freq_items_k, support_k = scan_dataset(dataset, next_candidates, min_support)
        support_data.update(support_k)
        freq_items.append(freq_items_k)
        k += 1

    return freq_items, support_data

def generate_rules(freq_items, support_data, min_confidence=0.7):
    # 生成关联规则
    rules = []
    for i in range(1, len(freq_items)):
        for freq_set in freq_items[i]:
            H1 = [frozenset([item]) for item in freq_set]
            if i > 1:
                rules_from_conseq(freq_set, H1, support_data, rules, min_confidence)
            else:
                calc_confidence(freq_set, H1, support_data, rules, min_confidence)

    return rules

def calc_confidence(freq_set, H, support_data, rules, min_confidence=0.7):
    # 计算规则的置信度
    pruned_H = []
    for conseq in H:
        conf = support_data[freq_set] / support_data[freq_set - conseq]
        if conf >= min_confidence:
            print(freq_set - conseq, '-->', conseq, 'conf:', conf)
            rules.append((freq_set - conseq, conseq, conf))
            pruned_H.append(conseq)

    return pruned_H

def rules_from_conseq(freq_set, H, support_data, rules, min_confidence=0.7):
    # 生成关联规则的递归函数
    m = len(H[0])
    if len(freq_set) > (m + 1):
        Hmp1 = apriori_gen(H, m + 1)
        Hmp1 = calc_confidence(freq_set, Hmp1, support_data, rules, min_confidence)
        if len(Hmp1) > 1:
            rules_from_conseq(freq_set, Hmp1, support_data, rules, min_confidence)

# 示例数据集
dataset = load_dataset()
# 设置最小支持度
min_support = 0.5
# 使用Apriori算法挖掘频繁项集
freq_items, support_data = apriori(dataset, min_support)
# 输出频繁项集
print("频繁项集：", freq_items)
# 输出支持度数据
print("支持度数据：", support_data)
# 设置最小置信度
min_confidence = 0.7
# 生成关联规则
rules = generate_rules(freq_items, support_data, min_confidence)
# 输出关联规则
print("关联规则：", rules)
